import os
from datetime import datetime

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import tensorflow as tf
from tensorflow import keras
import matplotlib.pyplot as plt
from tensorflow.keras import datasets, layers,optimizers, Sequential, metrics
import io
def preposs(x,y):
    '''数据预处理'''
    x = tf.cast(x, dtype=tf.float32)/255.
    y = tf.cast(y, dtype=tf.int32)

    return x,y

# 数据加载
(x,y),(x_test,y_test) = datasets.fashion_mnist.load_data()
print(x.shape,y.shape)

batchsz = 128 # 设置batch大小

db = tf.data.Dataset.from_tensor_slices((x,y))
db = db.map(preposs).shuffle(1000).batch(batchsz)#数据预处理、数据打散，分batch

db_test = tf.data.Dataset.from_tensor_slices((x_test,y_test))
db_test = db_test.map(preposs).shuffle(1000).batch(batchsz)#数据预处理、数据打散，分batch

db_iter = iter(db)
sample = next(db_iter)
print('batch:',sample[0].shape,sample[1].shape)


model = tf.keras.models.load_model('savedmodel/model.h5')
print('loadon model')
model.summary()

total_correct = 0
total_num = 0
for x, y in db_test:
    # x: [b,28,28] => [b,784]
    # y: [b]
    x = tf.reshape(x, [-1, 28 * 28])
    # [b,10]
    logits = model(x)
    prob = tf.nn.softmax(logits, axis=1)
    indices = tf.argmax(prob, axis=1)  # [b]
    indices = tf.cast(indices, dtype=tf.int32)

    result = tf.cast(tf.equal(indices, y), dtype=tf.int32)  # [b]
    correct = tf.reduce_sum(result)

    total_correct += int(correct)
    total_num += x.shape[0]

acc = total_correct / total_num

print('acc = ', acc)